import os
from openai import AsyncOpenAI
from agents import (
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    set_tracing_disabled,
)

from dotenv import load_dotenv

load_dotenv()

#disabled openai 模型调试模式
set_tracing_disabled(disabled=True)
class CustomProvider(ModelProvider):
    # 定义支持的 provider 名称列表，作为类属性，方便复用
    SUPPORTED_PROVIDERS = ['deepseek', 'moonshot', 'ollama', 'gitee', 'test']
    # 定义每个 provider 对应的环境变量前缀，作为类属性，方便复用
    PROVIDER_ENV_MAPPING = {
        'deepseek': 'DEEPSEEK',
        'moonshot': 'MOONSHOT',
        'ollama': 'OLLAMA',
        'gitee': 'GITEE',
        'test': 'TEST'
    }

    # init
    def __init__(self, provider_name: str) -> None:
        if provider_name not in self.SUPPORTED_PROVIDERS:
            raise ValueError(f"Unknown provider name: {provider_name}")
        
        self.provider_name = provider_name
        self.api_key = None
        self.base_url = None
        self.model_name = None

    def set_base_url(self, base_url: str):
        self.base_url = base_url

    def set_api_key(self, api_key: str):
        self.api_key = api_key

    def set_model(self, model_name: str):
        self.model_name = model_name

    def get_model(self, model_name: str = None) -> Model:
        if self.provider_name not in self.PROVIDER_ENV_MAPPING:
            raise ValueError(f"Unknown provider name: {self.provider_name}")

        env_prefix = self.PROVIDER_ENV_MAPPING[self.provider_name]

        # 确定模型名称，优先级：传入的 model_name > 实例属性 self.model_name > 环境变量
        model = model_name if model_name is not None else self.model_name or os.getenv(f"{env_prefix}_MODEL")
        # 获取 API 密钥，优先级：实例属性 self.api_key > 环境变量
        api_key = self.api_key or os.getenv(f"{env_prefix}_API_KEY")
        # 获取基础 URL，优先级：实例属性 self.base_url > 环境变量
        base_url = self.base_url or os.getenv(f"{env_prefix}_BASE_URL")

        # 增加参数检查，确保必要参数不为空
        if not model:
            raise ValueError(f"Model name is not provided for provider {self.provider_name}")
        if not api_key:
            raise ValueError(f"API key is not provided for provider {self.provider_name}")
        if not base_url:
            raise ValueError(f"Base URL is not provided for provider {self.provider_name}")

        self.base_url=base_url
        self.api_key=api_key
        self.model_name=model

        return OpenAIChatCompletionsModel(
            model=model,
            openai_client=AsyncOpenAI(
                api_key=api_key,
                base_url=base_url,
            ),
        )
